# 将 PaddleOCR-json box 改为(x,y,w,h) 结果可视化表现

from PIL import Image, ImageDraw, ImageFont


class visualize:
    """可视化"""

    # ================================ 静态方法 ================================

    def __init__(self, imagePath):
        """创建可视化对象。\n
        :textBlocks: 文本块列表，即OCR返回的data部分\n
        :imagePath: 对应的图片路径。
        """
        self.imgSource = Image.open(imagePath).convert("RGBA")  # 原始图片图层
        self.size = self.imgSource.size
        self.box_img = img = Image.new("RGBA", self.size, 0)
        self.box_draw = ImageDraw.Draw(img)
        #
        self.order_img = Image.new("RGBA", self.size, 0)
        self.order_draw = ImageDraw.Draw(img)
        ttfPath = "msyh.ttc"
        ttfSize = 17
        self.ttf = ImageFont.truetype(ttfPath, ttfSize)
        # self.imgBox = self.createBox(textBlocks, self.size)  # 包围盒图层
        # self.imgText = self.createText(textBlocks, self.size)  # 文字图层
        # self.imgOrder = self.createOrder(textBlocks, self.size)  # 序号图层

    def add_box(self, ocr_box, fill="#00500040", outline="#11ff22", width=2):
        """创建包围盒图层，返回PIL Image对象。\n
        :textBlocks: 文本块列表。\n
        :size: 图片尺寸。\n
        以下为可选字段：（颜色为十六进制6位RGB或8位RGBA字符串，如 #112233ff）\n
        :fill: 包围盒填充颜色。\n
        :outline: 包围盒轮廓颜色。\n
        :width: 包围盒轮廓粗细，像素。
        # [x, y], [x+w, y], [x, y+h], [x+w, y+h]
        # x y w, h
        """
        x, y, w, h = ocr_box
        box = [
            (x, y),
            (x + w, y),
            (x + w, y + h),
            (x, y + h),
        ]
        self.box_draw.polygon(box, fill=fill, outline=outline, width=width)

    def add_line_box(self, ocr_box, outline="#4682B4", width=2):
        """创建包围盒图层，返回PIL Image对象。\n
        :textBlocks: 文本块列表。\n
        :size: 图片尺寸。\n
        以下为可选字段：（颜色为十六进制6位RGB或8位RGBA字符串，如 #112233ff）\n
        :fill: 包围盒填充颜色。\n
        :outline: 包围盒轮廓颜色。\n
        :width: 包围盒轮廓粗细，像素。
        # [x, y], [x+w, y], [x, y+h], [x+w, y+h]
        # x y w, h
        """
        x, y, w, h = ocr_box
        box = [
            (x, y),
            (x + w, y),
            (x + w, y + h),
            (x, y + h),
        ]
        self.box_draw.polygon(box, outline=outline, width=width)

    def add_order(
            self,
            ocr_box,
            text,
            fill="#2233ff",
            bg="#ffffffe0",
    ):
        """创建序号图层，返回PIL Image对象。\n
        :textBlocks: 文本块列表。\n
        :size: 图片尺寸。\n
        以下为可选字段：\n
        :ttfPath: 字体文件路径。默认为微软雅黑，若不存在此字体会报错。\n
        :ttfSize: 字体大小。\n
        :fill: 文字颜色，十六进制6位RGB或8位RGBA字符串，如 #112233ff。\n
        """
        text = str(text)
        x, y, w, h = ocr_box

        xy = (x, y)  # 左上角坐标
        x_, y_, w, h = self.ttf.getbbox(text)  # 获取宽高。只需要w和h
        w *= 1.1
        h *= 1.1
        self.order_draw.rectangle((xy, (xy[0] + w, xy[1] + h)), fill=bg, width=0)  # 背景矩形
        self.order_draw.text(xy, text, font=self.ttf, fill=fill)  # 文字

    @staticmethod
    def createContrast(img1, img2):
        """左右拼合两个图片，创建对比图层，返回PIL Image对象。"""
        size = (img1.size[0] + img2.size[0], max(img1.size[1], img2.size[1]))
        print('--size--', size)
        img = Image.new("RGBA", size, 0)
        img.paste(img1, (0, 0))
        img.paste(img2, (img1.size[0], 0))
        return img

    @staticmethod
    def composite(img1, img2):
        """传入两个PIL Image对象（RGBA格式），以img1为底，将img2叠加在其上
        返回生成的图片"""
        return Image.alpha_composite(img1, img2)

    # ================================ 快捷接口 ================================

    def get(self, isBox=True, isText=False, isOrder=False, isSource=True):
        """返回合成可视化结果的PIL Image图像。\n
        :isBox: T时返回包围盒图层。\n
        :isText: T时返回文字图层。\n
        :isOrder: T时返回序号图层。\n
        :isSource: T时返回原图。F时返回透明背景的纯可视化结果。\n
        """
        img = Image.new("RGBA", self.size, 0)
        flags = (isSource, isBox, isText, isOrder)
        for index, im in enumerate(
                [self.imgSource, self.box_img, self.order_img]
        ):
            if im and flags[index]:
                img = visualize.composite(img, im)
        return img

    def show(self, isBox=True, isText=False, isOrder=False, isSource=True):
        """显示可视化结果图像。\n
        :isBox: T时返回包围盒图层。\n
        :isText: T时返回文字图层。\n
        :isOrder: T时返回序号图层。\n
        :isSource: T时返回原图。F时返回透明背景的纯可视化结果。\n
        """
        img = self.get(isBox, isText, isOrder, isSource)
        img.show()

    def save(self, path="", isBox=True, isText=False, isOrder=False, isSource=True):
        """保存可视化结果图像。\n
        :path: 保存路径。\n
        :isBox: T时返回包围盒图层。\n
        :isText: T时返回文字图层。\n
        :isOrder: T时返回序号图层。\n
        :isSource: T时返回原图。F时返回透明背景的纯可视化结果。\n
        """
        img = self.get(isBox, isText, isOrder, isSource)
        img.save(path)
